1918F - Caterpillar on a Tree - CodeForces Solution

#include<bits/stdc++.h>
using namespace std;
int read() {
	int f = 1, x = 0;
	char c = getchar();
	while (c < '0' || c > '9') {
		if (c == '-')f = -1;
		c = getchar();
	}
	while (c >= '0' && c <= '9') {
		x = x * 10 + c - '0';
		c = getchar();
	}
	return f * x;
}
void write(int x) {
	if (x < 0) {
		putchar('-');
		x = -x;
	}
	if (x > 9)write(x / 10);
	putchar(x % 10 + '0');
}
const int N = 2e5 + 10, MOD = 1e9 + 7, INF = 0x3f3f3f3f;
int fa[N], n = read(), k = read(), head[N], tot, dep[N], tmp[N], ans;
bool gg[N];
struct Edge {
	int to, nxt;
} e[N];
void add(int u, int v) {
	e[++tot].to = v;
	e[tot].nxt = head[u];
	head[u] = tot;
}
priority_queue<int>q;
void dfs(int pos) {
	dep[pos] = dep[fa[pos]] + 1;
	tmp[pos] = pos;
	for (int i = head[pos]; i; i = e[i].nxt) {
		dfs(e[i].to);
		if (dep[tmp[pos]] < dep[tmp[e[i].to]])tmp[pos] = tmp[e[i].to];
	}
}
void sfd(int pos) {
	for (int i = head[pos]; i; i = e[i].nxt) {
		sfd(e[i].to);
		if (tmp[pos] != tmp[e[i].to])q.push(-2 * dep[pos] + dep[tmp[e[i].to]]);
	}
}
signed main() {
	//freopen(".in", "r", stdin);
	//freopen(".out", "w", stdout);
	for (int i = 2; i <= n; i++) {
		fa[i] = read();
		add(fa[i], i);
	}
	dep[0] = -1;
	dfs(1);
	sfd(1);
	ans = 2 * n - 2 - dep[tmp[1]];
	while (!q.empty() && k) {
		if (q.top() <= 0)break;
		ans -= q.top();
		q.pop();
		k--;
	}
	write(ans);
	return 0;
}